require 'spec_helper'

RSpec.shared_examples "An AutoCheck module which can be overridden" do |opts|
  context "When the check method returns #{opts[:check_code]}", focus: opts[:focus] do
    let(:check_result) { opts[:check_code] }

    context "when the check method returns #{opts[:check_code]}" do
      context "when ForceExploit is enabled" do
        before(:each) do
          subject.datastore['ForceExploit'] = true
          subject.send(opts[:method])
        end

        it "calls the check method" do
          expect(subject).to have_received(:check)
        end

        it "calls the original #{opts[:method ]} method" do
          expect(subject).to have_received(:"original_#{opts[:method]}_call")
        end
      end
    end

    context "when ForceExploit is disabled" do
      before(:each) do
        subject.datastore['ForceExploit'] = false
      end

      it "it doesn't call the original #{opts[:method ]} method" do
        expect { subject.send(opts[:method]) }.to raise_error(opts[:expected_error]) do
          expect(subject).to_not have_received(:"original_#{opts[:method]}_call")
        end
      end
    end
  end
end

RSpec.shared_examples "An AutoChecked method" do |opts|
  subject { mock_module_with_prepend_autocheck.new }

  before(:each) do
    allow(subject).to receive(:check).and_return(check_result)
    allow(subject).to receive(:"original_#{opts[:method]}_call").and_call_original
  end

  context 'when AutoCheck is disabled' do
    let(:check_result) { ::Msf::Exploit::CheckCode::Vulnerable }

    before(:each) do
      subject.datastore['AutoCheck'] = false
      subject.send(opts[:method])
    end

    it "doesn't call the check method" do
      expect(subject).to_not have_received(:check)
    end

    it "correctly calls the #{opts[:method]} method" do
      expect(subject).to have_received(:"original_#{opts[:method]}_call")
    end
  end

  context 'when AutoCheck is enabled' do
    before(:each) do
      subject.datastore['AutoCheck'] = true
    end

    context 'when the check method returns vulnerable' do
      let(:check_result) { ::Msf::Exploit::CheckCode::Vulnerable }

      before(:each) do
        subject.send(opts[:method])
      end

      it "calls the check method" do
        expect(subject).to have_received(:check)
      end

      it "calls the original #{opts[:method]} method" do
        expect(subject).to have_received(:"original_#{opts[:method]}_call")
      end
    end

    context 'when the check method returns appears' do
      let(:check_result) { ::Msf::Exploit::CheckCode::Appears }

      before(:each) do
        subject.send(opts[:method])
      end

      it "calls the check method" do
        expect(subject).to have_received(:check)
      end

      it "calls the original #{opts[:method]} method" do
        expect(subject).to have_received(:"original_#{opts[:method]}_call")
      end
    end

    it_behaves_like "An AutoCheck module which can be overridden",
                    method: opts[:method],
                    check_code: ::Msf::Exploit::CheckCode::Safe,
                    expected_error: 'The target is not exploitable. "set ForceExploit true" to override check result.'

    it_behaves_like "An AutoCheck module which can be overridden",
                    method: opts[:method],
                    check_code: ::Msf::Exploit::CheckCode::Unsupported,
                    expected_error: 'This module does not support check. "set ForceExploit true" to override check result.'

    it_behaves_like "An AutoCheck module which can be overridden",
                    method: opts[:method],
                    check_code: ::Msf::Exploit::CheckCode::Unknown,
                    expected_error: 'Cannot reliably check exploitability. "set ForceExploit true" to override check result.'
  end
end

RSpec.describe Msf::Exploit::Remote::AutoCheck do
  let(:mock_module_with_prepend_autocheck) do
    context_described_class = described_class
    Class.new(::Msf::Exploit) do
      prepend context_described_class

      def check
        # mocked
      end

      def run
        original_run_call
      end

      def original_run_call
        # Helper for verifying the original run function was called
      end

      def exploit
        original_exploit_call
      end

      def original_exploit_call
        # Helper for verifying the original exploit function was called
      end
    end
  end
  let(:mock_module_with_include_autocheck) do
    context_described_class = described_class
    Class.new(::Msf::Exploit) do
      include context_described_class
    end
  end

  context 'When the mixin is included' do
    it "ensures that this scenario can not happen" do
      expect { mock_module_with_include_autocheck.new }.to raise_error NotImplementedError, "Msf::Exploit::Remote::AutoCheck should not be included, it should be prepended"
    end
  end

  context 'when the mixin is prepended' do
    describe '#run' do
      it_behaves_like 'An AutoChecked method', method: :run
    end

    describe '#exploit' do
      it_behaves_like 'An AutoChecked method', method: :exploit
    end
  end
end
